package com.shm.leetcode;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author: shm
 * @dateTime: 2020/10/6 9:08
 * @description: 834. 树中距离之和
 * 给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。
 *
 * 第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
 *
 * 返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
 *
 * 示例 1:
 *
 * 输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
 * 输出: [8,12,6,10,10,10]
 * 解释:
 * 如下为给定的树的示意图：
 *   0
 *  / \
 * 1   2
 *    /|\
 *   3 4 5
 *
 * 我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
 * 也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。
 * 说明: 1 <= N <= 10000
 */
public class SumOfDistancesInTree {

    List<List<Integer>> graph;
    int[] distSum;
    int[] nodeSum;
    int sum;
    /**
     * @author: shm
     * @dateTime: 2020/10/6 10:48
     * @description: 树中距离之和 | 树形DP | 详解
     * 思路
     *
     * 如上图，我们看节点 2，其他节点和它之间的距离和，可以分为两部分：
     *
     * 节点 2 所在的子树中的节点，与它的距离和。
     * 子树之外的节点与它的距离和。
     * 前者是一个子树内部的问题，我们寻找一下递归关系。
     *
     * 定义distSum[i]表示：节点i所在的子树中的节点到节点i的距离和。
     *
     * distSum[2]仅仅是它的子节点的distSum累加吗？不是的，还要算上走红色分支（如图）的次数。
     *
     * 节点 2 到子树 3 中的节点，要走两次红色路径，因为子树 3 有两个节点。
     * 节点 2 到子树 4 中的节点，要走一次红色路径，因为子树 4 有一个节点。
     * 节点 2 到子树 5 中的节点，要走一次红色路径，因为子树 5 有一个节点。
     *
     * 因此还需要计算每个子树的节点个数nodeNum[i]，递归公式如下：
     *
     * nodeNum[root] = sum(nodeNum[child])+1
     * nodeNum[root]=sum(nodeNum[child])+1
     *
     * distSum[root] = sum(nodeNum[child] + distSum[child])
     * distSum[root]=sum(nodeNum[child]+distSum[child])
     *
     * 从上往下递归，到达底部就遇到「结果已知的 base case」，随着递归的出栈，结果不断向上返回，最后求出当前子树的 distSum。
     *
     * 什么时候结束递归呢，即，什么时候才算遇到 base case 呢？
     *
     * 显然，我们需要遍历当前 root 节点的所有邻居，如果它的邻居只有它的父亲，说明它除了父亲，没有子节点，它就是 base case：nodeNum 为 1，distSum 为 0。
     *
     * 因为我们要遍历所有的邻居，所以需要根据 edges 数组，构建邻接关系表：
     *
     *
     * const graph = new Array(N);
     * for (let i = 0; i < graph.length; i++) {
     *     graph[i] = [];
     * }
     * for (const edge of edges) {
     *     const [from, to] = edge;
     *     graph[from].push(to);
     *     graph[to].push(from);
     * }
     * 这是 from bottom to top 的计算，由小的子树返回的结果，计算出大的子树的结果，即，在递归完当前节点的所有子树之后，再执行累加计算当前节点。在递归出栈时解决问题。
     *
     * 因此，采用后序遍历：
     * const postOrder = (root, parent) => {
     *     const neighbors = graph[root]; // 与它相连的节点们
     *     for (const neighbor of neighbors) {
     *         if (neighbor == parent) {  // 如果邻居是自己父亲，跳过。
     *             continue;              // 如果邻居只有自己父亲，则for循环结束，当前递归结束
     *         }
     *         postOrder(neighbor, root);
     *         nodeNum[root] += nodeNum[neighbor];
     *         distSum[root] += nodeNum[neighbor] + distSum[neighbor];
     *     }
     * };
     * ---------------------------------------------- 我是分隔符 ----------------------------------------------------
     * 此时的distSum数组，存放 root 到所在子树中的节点的距离和，显然，还要加上子树外的节点到 root 的距离和，才是最后的结果。
     *
     * 后者怎么求？不好求。我们需要求吗？其实不需要。
     *
     * 我们看到，对整个树的根节点来说，它的distSum已经是对的，因为，它的子树就是整棵树。
     *
     * 我们这样：根据已知的对的distSum，更新它下面的子节点的distSum。from top to bottom.
     *
     * 更新计算子节点i的distSum[i]，怎么用上已知的distSum[root]呢？
     *
     *
     *
     * 如上图，节点2所在的子树的节点个数为nodeNum[2]，这nodeNum[2]个节点，从计算distSum[0]变成计算distSum[2]：从节点 0 到这些节点，变成从节点 2 到这些节点，每个节点都少走了一步，一共少走了nodeNum[2]步。
     *
     * 子树以外的节点呢，有N - nodeNum[2]个，从计算distSum[0]变成计算distSum[2]：从节点 0 到这些节点，变成从节点 2 到这些节点，每个节点都多走了一步，一共多走了N-nodeNum[2]步。
     *
     * 所以 distSum[i] 的更新公式如下：
     *
     * distSum[i] = distSum[root] - nodeNum[i] + (N - nodeNum[i])
     * distSum[i]=distSum[root]−nodeNum[i]+(N−nodeNum[i])
     *
     * 这是从上往下更新，不要等到遇到底部的 base case 向上返回，就要计算。即，在压栈前解决问题：在递归当前节点的子树之前，就更新出当前节点的distSum，使之正确，在后续的递归中，它作为正确的父节点 distSum，递推出子节点的 distSum。
     *
     * 因此采用前序遍历：
     *
     *
     * const preOrder = (root, parent) => {
     *     const neighbors = graph[root];
     *     for (const neighbor of neighbors) {
     *       if (neighbor == parent) {
     *         continue;
     *       }
     *       distSum[neighbor] = distSum[root] - nodeNum[neighbor] + (N - nodeNum[neighbor]);
     *       preOrder(neighbor, root);
     *     }
     *   };
     * 可以看到，我们用了两次 DFS 搜索，前者是计算出每个子树的 nodeNum 和 distSum，后者是计算出真正的 distSum，即每个节点与「其他所有节点」的距离和，注意二者的区别。
     *
     * 复盘总结
     * 我们重复利用了distSum数组这个空间。它先是代表：节点到它所在子树的节点的距离和。后面就更新为：节点到其他所有节点的距离和。
     *
     * 当然，我们也可以开辟一个新的容器，毕竟它们两个是不同的定义。但没必要，节省一点空间。
     *
     * 本题的突破口，依然是——总是关注子树，关注子树，想到递归，子树是一种重复的结构，适合用递归求解，这道题巧妙地将目标拆分为两部分：
     *
     * 子树内的节点与 root 的距离和
     * 子树外的节点与 root 的距离和
     * 用递归求解前者，这是主旋律，后者不用特地求，因为我们已知父节点的正确的 distSum，利用递推公式：
     *
     * distSum[i] = (distSum[root] - nodeNum[i]) + (N - nodeNum[i])
     * distSum[i]=(distSum[root]−nodeNum[i])+(N−nodeNum[i])
     *
     * 就能递推出当前节点的正确的 distSum，就像沿着一棵树在填表，这就是树形DP。
     * 作者：xiao_ben_zhu
     * 链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree/solution/shou-hua-tu-jie-shu-zhong-ju-chi-zhi-he-shu-xing-d/
     */
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        sum =N;
        // 建立映射表，graph[i]：存放与节点i相连的所有节点
        graph = new ArrayList<>();
        for (int i = 0; i < N; i++) {
            graph.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            graph.get(edge[0]).add(edge[1]);
            graph.get(edge[1]).add(edge[0]);
        }

        // distSum[i]：节点i到它所在子树的节点的距离和，后面更新为：节点i到其他所有节点的距离和
        distSum = new int[N];
        nodeSum = new int[N];
        // nodeNum[i]：节点i所在子树的节点个数，保底为1
        Arrays.fill(nodeSum,1);
        postDfs(0,-1);
        preDfs(0,-1);
        return distSum;
    }

    public void postDfs(int root,int parent){
        List<Integer> neighbors = graph.get(root);
        for (Integer neighbor : neighbors) {
            // 如果邻居是自己父亲，跳过。
            // 如果邻居只有自己父亲，则for循环结束，当前递归结束
            if (neighbor==parent){
                continue;
            }
            postDfs(neighbor,root);
            nodeSum[root]+=nodeSum[neighbor];
            distSum[root]+=nodeSum[neighbor]+distSum[neighbor];
        }
    }

    public void preDfs(int root,int parent){
        List<Integer> neighbors = graph.get(root);
        for (Integer neighbor : neighbors) {
            if (neighbor==parent){
                continue;
            }
            distSum[neighbor]=distSum[root]-nodeSum[neighbor]+(sum-nodeSum[neighbor]);
            preDfs(neighbor,root);
        }
    }

}
